# 把合成的大文件进行频率成分分析

# pre packages
from myglobal import os, sys, LINES

# sys packages
from pylab import np
import obspy as ob
from glob import glob
from tqdm import tqdm
from obspy.core.utcdatetime import UTCDateTime
import pandas as pd
from scipy.stats import linregress
import datetime
from scipy import interpolate


# self packages
from utils.loc import load_loc, get_distance,sort_data_by_distance
from utils.math import measure_shift_fft,remove_point_skip,analyze_frequency_components
from utils.h5data import get_event_data, save_h5_data, read_h5_data,h5glob
from utils.plot import plot_scatters,plot_traces
from utils.trace import get_tapered_slices, safe_filter,get_tapered_traces


def plot_all_freqs(te,data_all, ):

    fs=0.3
    fe=4

    dte = te[1]-te[0]
    STATS = data_all.keys()
    dt_all = np.array([data_all[name][0] for name in STATS])
    T_MAX_all = np.array([data_all[name][1] for name in STATS])
    x = np.array([data_all[name][2] for name in STATS])
    NX = len(x)
    P_all = []
    for j, name in enumerate(STATS):
        dt_j,T_MAX,xj, az_j, health_j = data_all[name]
        dt_j[health_j==0]=0
        dt_j[te>90]=0
        # freqs, P, amp, angle = get_spec(dt_j, dte,fs,fe)
        freqs, P, amp, angle = get_spec_lombscargle(dt_j, dte,fs,fe)
        # P = P/(T_MAX+0.00001)
        print(name, T_MAX)
        P_all.append(P)
    P_all = np.array(P_all)
    CP1 = np.argmin(np.abs(freqs-1))
    CP2 = np.argmin(np.abs(freqs-2))

    from pylab import figure, plt
    from matplotlib import rcParams, gridspec
    rcParams['font.family'] = 'Arial'
    rcParams['font.size'] = '8'
    DPI=600

    SCALE = 1/P_all.var()**0.5 *100
    XLIM=[0,20]
    YLIM =[x.min()-400,x.max()+400]

    plt.close('all')
    fig = figure(figsize=[6,3], dpi=DPI)
    gs = fig.add_gridspec(1,4,wspace=0.1)

    ax1 = fig.add_subplot(gs[0,0:2])
    summary = 0
    for j, name in enumerate(STATS):
        amp = np.abs(P_all[j,:])
        summary+=amp/amp.max()
        ax1.plot(freqs, amp*SCALE+x[j], '-', color='k', lw=0.7, label=f'{name}')
        ax1.text(fe-0.5,x[j], name, fontsize=8)

    ax1.plot(freqs, summary*SCALE+YLIM[0]+100, '-', color='tab:red', lw=0.7, label=f'{name}')
    ax1.text(fe-0.5,YLIM[0]+100, 'sum', fontsize=8)

    # ax1.scatter([1]*len(x),np.abs(P_all[:, CP1])*SCALE+x, s=10, c='tab:red', marker='^', edgecolors='none')
    # ax1.scatter([2]*len(x),np.abs(P_all[:, CP2])*SCALE+x, s=10, c='tab:blue', marker='o', edgecolors='none')
    ax1.plot([freqs[CP1],freqs[CP1]],YLIM,c='tab:red', lw=0.7)
    ax1.plot([freqs[CP2],freqs[CP2]],YLIM,c='tab:blue', lw=0.7)
    ax1.plot([1.93,1.93],YLIM,c='tab:blue', lw=0.7)


    ax1.set_xlabel('freq(cycle/day)')
    ax1.set_ylabel('x(m)')
    ax1.set_ylim(YLIM)
    # ax1.set_xlim([1.9,2.1])
    ax1.set_xlim([fs,fe])
    ax1.set_title('spectrum')
    
    ax2 = fig.add_subplot(gs[0,2])

    ax2.scatter(np.angle(P_all[:, CP1])/np.pi, x, s=10, c='tab:red', marker='^', edgecolors='none')
    ax2.scatter(np.angle(P_all[:, CP2])/np.pi, x, s=10, c='tab:blue', marker='o', edgecolors='none')

    ax2.set_xlabel(r'$\phi$($\pi$)')
    ax2.set_ylim(YLIM)
    ax2.set_xlim([-1,1])
    ax2.set_xticklabels(['-$\pi$', '0', '$\pi$'])
    ax2.set_yticks([])
    ax2.set_title('phase Vs r')

    ax3 = fig.add_subplot(gs[0,3])

    ax3.scatter(np.abs(P_all[:, CP1]), x, s=10, c='tab:red', marker='^', edgecolors='none')
    ax3.scatter(np.abs(P_all[:, CP2]), x, s=10, c='tab:blue', marker='o', edgecolors='none')
    ax3.set_xlabel('A(ms)')
    ax3.set_ylim(YLIM)
    MAX1 = 3*np.abs(P_all[:, CP1]).var()**0.5
    MAX2 = 3*np.abs(P_all[:, CP2]).var()**0.5
    ax3.set_xlim([0,max(MAX1,MAX2)])
    # ax3.set_xlim([-20,20])
    ax3.set_yticks([])
    ax3.set_title('amp Vs r')
    
    return fig

def get_spec(x,dt, fs,fe):
    from scipy import signal

    nx = len(x)
    y1 = get_tapered_traces(x,dt,L_Taper=dt*nx/20)

    freqs = np.fft.fftfreq(nx, dt)
    P = np.fft.fft(y1)

    fsp = np.argmin(np.abs(freqs-fs))
    fep = np.argmin(np.abs(freqs-fe))
    freqs = freqs[fsp:fep]
    P = P[fsp:fep]

    N_VALID = (x!=0).sum()+1
    return freqs,P/N_VALID,np.abs(P)/N_VALID, np.angle(P)

def get_spec_lombscargle(x,dt, fs,fe):
    from scipy import signal

    nx = len(x)
    t = np.arange(nx)*dt
    freqs = np.arange(fs,fe,0.002)
    freqs, P, A, phi= analyze_frequency_components(t,x, freqs)
    return freqs,P,A, phi


def get_data_all_from_8file(INPUT,STATS, key):
    data_all = {}
    for j, name in enumerate(STATS):
        dt_j = read_h5_data(INPUT,[f'{key}'] ,group_name=f'{name}/merged')[0]
        health = read_h5_data(INPUT,[f'{key}'] ,group_name=f'{name}/health')[0]
        T_MAX,r_j, az_j = read_h5_data(INPUT,['T_MAX','r','az'] ,group_name=f'{name}/RAW/{key}')
        # if r_j>100:
        #     continue
        if abs(r_j)>1500:
            continue
        data_all[name]=[dt_j,T_MAX,-r_j, az_j,health]
    return data_all
def get_data_phy_from_8file(INPUT,keys):
    data_all = {}
    for j, name in enumerate(keys):
        phy_j = read_h5_data(INPUT,[f'{name}'] ,group_name=f'/phy')[0]
        health = np.ones_like(phy_j)
        data_all[name]=[phy_j,0,0, 0,health]
    return data_all

if __name__ == '__main__':
# cmd
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-debug', action='store_true', help='method: DEBUG')
    parser.add_argument('-figroot', default='figures/8.merge.dtshift', help='root to save figs')
    parser.add_argument('-keyN', type=int, default=-1,  help='which key.idx to plot')
    parser.add_argument('-input', default='',  help='file pattern to load')
    args = parser.parse_args()
    print(args)

    DEBUG= args.debug
    INPUT = args.input
    FIG_ROOT = args.figroot
    keyN = args.keyN
    # 读取原始数据
    datasets,args_infile= read_h5_data(INPUT,['all_groups','keys','files','te'], group_name='metadata', read_attrs=True)
    STATS = [i.decode() for i in datasets[0]] #台站名
    keys = [i.decode() for i in datasets[1]] # 数据名
    files = [i.decode() for i in datasets[2]] 
    te = datasets[3]
    dte = te[1]-te[0]

    # 处理元数据
    print(keyN)
    DATE=args_infile['date']
    EMARKER = args_infile['emarker']

    FIG_ROOT = f'{FIG_ROOT}/8BC.{DATE}.{EMARKER}'
    if not os.path.exists(FIG_ROOT):
        os.makedirs(FIG_ROOT)

    # 物理信号
    # phy_keys = h5glob(INPUT, '*',group_name='/phy',object_type='datasets', STRIP_GROUP_NAME=True)
    # phy_items = read_h5_data(INPUT, phy_keys, group_name='/phy')
    # data_phy = {}
    # for i in range(len(phy_keys)):
    #     data_phy[phy_keys[i]] = phy_items[i]
    keys_used = keys if keyN==-1 else [keys[keyN]]
    for key in keys_used:
        data_all_key = get_data_all_from_8file(INPUT,STATS, key)
        fig = plot_all_freqs(te,data_all_key)
        fig.tight_layout()
        output_fig = f'{FIG_ROOT}/spec_{key}.png'
        print(output_fig,)
        fig.savefig(output_fig)
    phy_keys = h5glob(INPUT, '*',group_name='/phy',object_type='datasets', STRIP_GROUP_NAME=True)
    for i, key in enumerate(phy_keys):
        data_all_key = get_data_phy_from_8file(INPUT, [key])
        fig = plot_all_freqs(te,data_all_key)
        fig.tight_layout()
        output_fig = f'{FIG_ROOT}/phy_spec_{key}.png'
        print(output_fig,)
        fig.savefig(output_fig)
